/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: op remote look up execute.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#ifndef OP_REMOTE_LOOK_UP_EXECUTOR_H
#define OP_REMOTE_LOOK_UP_EXECUTOR_H

#include <map>
#include <unordered_map>
#include <atomic>
#include "op_executor.h"
#include "transport_heterog_pub.h"
#include "hccl_operator.h"
#include "network_manager_pub.h"
#include "mr_manager.h"
#include "auto_buffer.h"
#include "transport_heterog_pub.h"
#include "profiling_manager.h"

#ifndef HCCD
#include "tbe_gatherv2_aicore.h"
#endif

namespace hccl {

using KeyPosRecord = struct KeyPosRecordDef {
    u32 psId;
    u64 pos;
    KeyPosRecordDef() : psId(0), pos(0) {}
};

class OpRemoteLookUpExecutor : public OpExecutor {
public:
    explicit OpRemoteLookUpExecutor(const std::unique_ptr<MrManager> &mrManager,
        const HcclDispatcher dispatcher, const u32 localWorldRank);
    ~OpRemoteLookUpExecutor() override;
    HcclResult Init(std::map<u32, std::unique_ptr<TransportHeterog>> &transportHandleMap,
        std::map<u32, u32> &psRankIdMap);
    HcclResult ExecuteOp(HcclOpPtr opPtr) override;

private:
    struct DropDuplicatesParams {
        PartitionMapTask task;
        OpRemoteLookUpExecutor *selpPtr;
    };

    struct RankIdAndPos {
        u32 pos{};
        u32 rankId{};
    };

    static void DropDuplicatesCallback(void *fnData);
    static void PrintValue(void *fnData);

    HcclResult TaskSend(OpTaskPtr &task);
    HcclResult TaskRecv(OpTaskPtr &task);
    HcclResult TaskIsend(OpTaskPtr &task);
    HcclResult TaskWait(OpTaskPtr &task);

    HcclResult TaskMemcpy(OpTaskPtr &task);

    HcclResult TaskDropDuplicates(OpTaskPtr &task);
    HcclResult TaskRemoteLookupProcess(OpTaskPtr &task);

    HcclResult TaskSendKeysByRdma(OpTaskPtr &task);
    HcclResult PrepareHostEnvelopeTemplates(OpTaskPtr &task);
    HcclResult FillEnvelope(OpTaskPtr &task, u32 rankId, HcclEnvelope &envelope);
    HcclResult AsyncCopyHostEnvelopeToDev(OpTaskPtr &task);

    HcclResult WaitSendKeyFinish(OpTaskPtr &task) const;

    HcclResult DeInit();

    HcclResult ExecuteTask(HcclOperator<EmbeddingServiceParam> *opPtr);
    HcclResult TaskImprobeRecv(OpTaskPtr &task);

    TransportHandle GetTransport(TransTask *taskPtr);

    // 推理lookup接口
    HcclResult KeysSort(DropDuplicatesParams &params,
        std::map<u32, u32> &rankIdKeyNumMap);
    HcclResult SendKeys(DropDuplicatesParams &params, std::map<u32, u32> &rankIdKeyNumMap,
        std::map<u32, u64> &waitLookUpRanks);
    HcclResult RecvValue(DropDuplicatesParams &params, std::map<u32, u64> &waitLookUpRanks);
    HcclResult RecoverValue(DropDuplicatesParams &params) const;

    HcclResult TaskKeyDropDuplicatesKernelLaunch(OpTaskPtr &task);
    HcclResult TaskSendKeysKernelLaunch(OpTaskPtr &task);
    HcclResult TaskSendKeysRecord(OpTaskPtr &task);
    HcclResult TaskEventWaitRecvDone(OpTaskPtr &task);
    HcclResult TaskResetUniqueHandleKernelLaunch(OpTaskPtr &task);
    HcclResult TaskEventWaitSendDone(OpTaskPtr &task);
    HcclResult TaskRecvValuesKernelLaunch(OpTaskPtr &task);
    HcclResult TaskRecoverValueKernelLaunch(OpTaskPtr &task);
    HcclResult TaskGatherFinishKernelLaunch(OpTaskPtr &task);
    HcclResult PrepareMappingMatrixForGatherValue(DropDuplicatesParams &params);

    HcclResult RecoverValueAicoreKernelLaunch(OpTaskPtr &task);
    HcclResult InferRecoverValueAicoreKernelLaunch(const PartitionMapTask &mapTask);
    HcclResult ChkServiceCancelKernelLaunch(const PartitionMapTask &mapTask);
    u32 GetTimeOut();
    void SetupTaskHdcsEmbeddingServiceParam(PartitionMapTask *taskPtr, HdcsEmbeddingServiceParam &para) const;
    HcclResult ProfReportTransportShmEvent(TransportHeterog *transPtr, uint64_t beginTime) const;
    EsLoopUpPara CreateEsLoopUpPara(const PartitionMapTask* taskPtr) const;
    
    std::map<OpTaskType, std::function<HcclResult(OpTaskPtr &)>> taskExeFuncTable_;
    NetworkManager *netResource_ = nullptr;
    std::unordered_map<u32, TransportHeterog *> rankTransportMap_;
    std::queue<TransportHandle> tempTransQue_;

    DeviceMem defaultSendMem_;
    DeviceMem defaultRecvMem_;
    u32 *psIdsMem_ = nullptr;
    u32 keyMaxNum_ = 0;
    std::vector<u32> psRankId_;
    u32 psSize_ = 0;
    std::vector<u64> psIdOffset_;
    std::vector<u32> rankId2PsId_;
    std::atomic<bool> blockFlag_ = {true};
    void *errorMsgAddr_ = nullptr;
    const std::unique_ptr<MrManager> &mrManager_;
    const HcclDispatcher dispatcher_;
    const u32 localWorldRank_;

    std::vector<KeyPosRecord> keyPosRecord_;
    std::vector<void *> psValueStartAddr_;
    std::queue<DropDuplicatesParams> dropDuplicatesQue_;

    AutoBuffer<s64> mappingMatrixS64_{};
    AutoBuffer<s32> mappingMatrix_{};
    std::unordered_map<u32, u64> rankIdAccumOffset_{};

    std::unique_ptr<HcclEsRdmaInfoForLookup[]> rdmaEnvelopeInfoHostMem_{};

#ifndef HCCD
    TbeReduce::TbeGatherV2 aiCoreGatherV2_{};
#endif
};
}
#endif
